use std::future::Future;
use std::pin::Pin;
use std::time::Duration;
use crate::context::Context;
use crate::message::DynamicMessage;
use crate::errors::SendError;
use crate::middleware::{Middleware, Next};

pub struct ExponentialBackoff {
    initial_delay: Duration,
    max_delay: Duration,
    multiplier: f64,
}

impl ExponentialBackoff {
    pub fn new(initial_delay: Duration, max_delay: Duration, multiplier: f64) -> Self {
        Self {
            initial_delay,
            max_delay,
            multiplier,
        }
    }

    pub fn get_delay(&self, attempt: usize) -> Duration {
        let delay = self.initial_delay.as_secs_f64() * self.multiplier.powi(attempt as i32);
        Duration::from_secs_f64(delay.min(self.max_delay.as_secs_f64()))
    }
}

impl Default for ExponentialBackoff {
    fn default() -> Self {
        Self {
            initial_delay: Duration::from_millis(100),
            max_delay: Duration::from_secs(30),
            multiplier: 2.0,
        }
    }
}

pub struct RetryMiddleware {
    max_retries: usize,
    backoff: ExponentialBackoff,
}

impl RetryMiddleware {
    pub fn new(max_retries: usize, backoff: ExponentialBackoff) -> Self {
        Self {
            max_retries,
            backoff,
        }
    }

    pub fn with_default_backoff(max_retries: usize) -> Self {
        Self {
            max_retries,
            backoff: ExponentialBackoff::default(),
        }
    }
}

impl Middleware for RetryMiddleware {
    fn handle<'a>(&'a self, ctx: &'a mut Context, msg: DynamicMessage, next: Next<'a>) -> Pin<Box<dyn Future<Output = Result<(), SendError>> + Send + 'a>> {
        Box::pin(async move {
            let mut retries = 0;
            loop {
                match next().await {
                    Ok(_) => return Ok(()),
                    Err(e) if retries < self.max_retries => {
                        retries += 1;
                        let delay = self.backoff.get_delay(retries);
                        tokio::time::sleep(delay).await;
                        continue;
                    }
                    Err(e) => return Err(e),
                }
            }
        })
    }
} 